package test.concurrent;

/*
 * Copyright 2013 (raistlic@gmail.com)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantLock;

/**
 * This class allows multiple threads to lock against a key of the specified type, 
 * based on the {@code equals()} contract defined in {@link java.lang.Object} 
 * class, in a reentrant manner.
 * 
 * <p/>
 * All methods defined in this class, including the static factory methods, are 
 * thread safe.
 * 
 * @param <K> the specified key type
 * 
 * @author icebamboo_moyun, raistlic
 */
public abstract class LockMap<K> {
  
  /**
   * This method creates and exports a {@code LockMap} instance that is "Garbage
   * Collector friendly", that is, it releases any lock references that is no 
   * longer used. It is recommended that user code with a large dictionary of 
   * keys favors this one over {@link #forSmallKeySet()}.
   * 
   * @param <K> the specified key type
   * 
   * @return the created {@code LockMap} instance.
   */
  public static <K> LockMap<K> forLargeKeySet() {
    
    return new GCedLockMap<K>();
  }
  
  /**
   * This method creates and exports a {@code LockMap} instance that reuses lock
   * objects against a key, note that the locks are not released until the 
   * {@code KeyMap} instance itself is GCed. It is recommended that user code 
   * with a small set of keys favors this one over {@link #forLargeKeySet()}.
   * 
   * @param <K> the specified key type
   * @return the created {@code LockMap} instance.
   */
  public static <K> LockMap<K> forSmallKeySet() {
    
    return new NotGCedLockMap<K>();
  }
  
  /**
   * This method helps the current calling thread to lock against the specified
   * {@code key}, or increase the lock value by 1, if it is already locked by the
   * current thread (reentrant lock).
   * 
   * @param key the specified key to lock, or increase lock value.
   * 
   * @throws InterruptedException if the current thread is interrupted while it 
   *         is waiting for a lock, the {@code LockMap} instance is guaranteed
   *         still in a stable state (as before the method is called) after the 
   *         exception is thrown.
   * 
   * @throws NullPointerException if the specified {@code key} is {@code null}.
   */
  public abstract void lock(K key) throws InterruptedException ;
  
  /**
   * This method decreases the lock value against the specified {@code key} by 1; 
   * and if the lock value hits 0 as a result of the decrement, the current 
   * calling thread unlocks the {@code key}.
   * 
   * @param key the specified key to decrease lock value, or unlock.
   * 
   * @throws NullPointerException if the specified {@code key} is {@code null}.
   * 
   * @throws IllegalArgumentException if the specified {@code key} is NOT 
   *         currently locked by the current thread.
   */
  public abstract void unlock(K key);
  
  /**
   * This method helps the current calling thread to check if it holds the lock
   * against the specified {@code key}.
   * 
   * @param key the specified key to check.
   * 
   * @return {@code true} if the current calling thread holds the lock against
   *         the specified {@code key}.
   * 
   * @throws NullPointerException if the specified {@code key} is {@code null}.
   */
  public abstract boolean isLockedByCurrentThread(K key);
  
  /*
   * This class is currently designed not to be inherited or directly 
   * instantiated from outside the file.
   */
  private LockMap() {}
  
  private static class GCedLockMap<K> extends LockMap<K> {
    
    private final ConcurrentHashMap<K, ReentrantLock> sharedLocks;
    
    private final ThreadLocal<Map<K, ReentrantLock>> localLocks; 
    
    private GCedLockMap() {
      
      sharedLocks = new ConcurrentHashMap<K, ReentrantLock>();
      localLocks = new ThreadLocal<Map<K, ReentrantLock>>() {
        
        @Override
        protected Map<K, ReentrantLock> initialValue() {
          
          return new HashMap<K, ReentrantLock>();
        }
      };
    }

    @Override
    public void lock(K key) throws InterruptedException {
      
      if( key == null )
        throw new NullPointerException();
      
      Map<K, ReentrantLock> localMap = localLocks.get();
      ReentrantLock localLock = localMap.get(key);
      if( localLock == null ) {
        
        localLock = new ReentrantLock();
        
        // potential failure point: safe to fail, won't effect any internal state.
        localLock.lockInterruptibly();
        localMap.put(key, localLock);
        
        ReentrantLock sharedLock = sharedLocks.put(key, localLock);
        if( sharedLock != null ) {
          
          // while waiting for the sharedLock, the current thread simply can not
          // be interrupted, or we're screwed.
          
          sharedLock.lock();
          sharedLock.unlock();
          
          // sharedLock (the old one, created by another thread) will then be GCed,
          // once its reference in its creating thread's ThreadLocal map is released.
        }
      }
      else {
        
        localLock.lockInterruptibly();
      }
    }

    @Override
    public void unlock(K key) {
      
      if( key == null )
        throw new NullPointerException();
      
      Map<K, ReentrantLock> localMap = localLocks.get();
      ReentrantLock localLock = localMap.get(key);
      if( localLock == null ) {
        
        throw new IllegalArgumentException(
                "Cannot release lock, because current thread does not hold "
                + "lock for the key: " + key);
      }
      
      int count = localLock.getHoldCount();
      if( count == 1 ) {
        
        sharedLocks.remove(key, localLock);
        localMap.remove(key);
      }
      localLock.unlock();
    }
    
    @Override
    public boolean isLockedByCurrentThread(K key) {
      
      if( key == null )
        throw new NullPointerException();
      
      return localLocks.get().containsKey(key);
    }
  }
  
  private static class NotGCedLockMap<K> extends LockMap<K> {
    
    private final Map<K, ReentrantLock> locks;
    
    private NotGCedLockMap() {
      
      locks = new ConcurrentHashMap<K, ReentrantLock>();
    }

    @Override
    public void lock(K key) throws InterruptedException {
      
      if( key == null )
        throw new NullPointerException();
      
      ReentrantLock lock = locks.get(key);
      
      if( lock == null ) synchronized (locks) {
        
        lock = locks.get(key);
        if( lock == null ) {
          
          lock = new ReentrantLock();
          locks.put(key, lock);
        }
      }
      lock.lockInterruptibly();
    }

    @Override
    public void unlock(K key) {
      
      if( key == null )
        throw new NullPointerException();
      
      ReentrantLock lock = locks.get(key);
      if( lock == null ) {
        
        throw new IllegalArgumentException(
                "Cannot release lock, because current thread does not hold "
                + "lock for the key: " + key);
      }
      lock.unlock();
    }
    
    @Override
    public boolean isLockedByCurrentThread(K key) {
      
      if( key == null )
        throw new NullPointerException();
      
      ReentrantLock lock = locks.get(key);
      return lock != null && lock.isHeldByCurrentThread();
    }
  }
}